#include "../include/rkmedia.h"
#include "../include/model.h"
#include <cstdint>
#include <cstring>
#include <opencv2/core/types.hpp>
#include <opencv2/imgproc.hpp>
#include "../include/rtsp_demo.h"
#include <sched.h>
#include <unistd.h>
#include <vector>

rtsp_demo_handle g_rtsplive = NULL;
static rtsp_session_handle g_rtsp_session = NULL;
rknn_ctx ctx;

void venc_rtsp(MEDIA_BUFFER mb)
{
    static RK_S32 packet_cnt = 0;

    if (g_rtsplive && g_rtsp_session) {
        rtsp_tx_video(g_rtsp_session, (uint8_t*)RK_MPI_MB_GetPtr(mb), RK_MPI_MB_GetSize(mb),
                  RK_MPI_MB_GetTimestamp(mb));
        rtsp_do_event(g_rtsplive);
        printf("#Get packet-%d, size %zu\n", packet_cnt, RK_MPI_MB_GetSize(mb));
    }

    RK_MPI_MB_ReleaseBuffer(mb);	//释放帧数据
    packet_cnt++;
}

int main()
{
    // 初始化
    VI_CHN_ATTR_S vi0_chn_attr;
    VENC_CHN_ATTR_S venc_chn_attr;
    RGA_ATTR_S stRgaAttr;
    MPP_CHN_S stSrcChn0;
    MPP_CHN_S stDestChn0;
    //output vector
    std::vector<float> filterBoxes;
    std::vector<float> objProbs;
    std::vector<int> classId; 
    std::vector<int> indexArray;
    int validCount;
    RK_MPI_SYS_Init();
    VI_init(&vi0_chn_attr, 0);
    init_rga(&stRgaAttr);
    VENC_init(&venc_chn_attr);
    VI_RGA_Bind(&stSrcChn0, &stDestChn0, 0);
    model_init(&ctx);
    /*------------rtsp初始化-------------*/
    g_rtsplive = create_rtsp_demo(8554);	//设置推流端口
    //推流的路径,注意这里的路径直接传字符串或者宏，不要传数组，会报错！！！
    g_rtsp_session= rtsp_new_session(g_rtsplive, "/live/main_stream");
    rtsp_set_video(g_rtsp_session, RTSP_CODEC_ID_VIDEO_H264, NULL, 0);
    //同步时间
	rtsp_sync_video_ts(g_rtsp_session, rtsp_get_reltime(), rtsp_get_ntptime());
    RK_MPI_VI_StartStream(0, 0);
    /*------------VENC数据输出-------------*/
    MPP_CHN_S stEncChn; // 声明模块设备通道
    stEncChn.enModId = RK_ID_VENC; // 指定模块号为VENC模块
    stEncChn.s32DevId = 0; // 设备号
    stEncChn.s32ChnId = 0; // 通道号
    // 注册数据输出回调函数
    RK_MPI_SYS_RegisterOutCb(&stEncChn, venc_rtsp);
    
    while(1)
    {
        MEDIA_BUFFER mb = RK_MPI_SYS_GetMediaBuffer(RK_ID_RGA, 0, -1);
        cv::Mat img(video_width, video_height, CV_8UC3, RK_MPI_MB_GetPtr(mb));
        RK_MPI_SYS_SendMediaBuffer(RK_ID_RGA, 1, mb);
        MEDIA_BUFFER MB = RK_MPI_SYS_GetMediaBuffer(RK_ID_RGA, 1, -1);
        // 获取推理结果     
        validCount = model_inference(&ctx, RK_MPI_MB_GetPtr(MB), filterBoxes, objProbs, classId, indexArray);
        // 画框
        printf("\n---------------------------\n");
        for(int i = 0; i < validCount; i++)
        {
            if(indexArray[i] != -1)
            {
                int x = filterBoxes[i*4]*video_width/640;
                int y = filterBoxes[4*i+1]*video_height/640;
                int w = filterBoxes[4*i+2]*video_width/640;
                int h = filterBoxes[4*i+3]*video_height/640;
                x = std::max(0, std::min(x, video_width - 1));
                y = std::max(0, std::min(y, video_height - 1));
                w = std::max(1, std::min(w, video_width - x));
                h = std::max(1, std::min(h, video_height - y));
                cv::Rect rect(x, y, w, h);
                cv::Point point(x, y);
                cv::rectangle(img, rect, cv::Scalar(0, 0, 255), 5);
                cv::putText(img, coco_class[classId[i]], point, cv::FONT_HERSHEY_PLAIN, 2, cv::Scalar(255, 0, 0), 2);
                printf("%s:(%d,%d,%d,%d)\n",coco_class[classId[i]].c_str(), x, y, w, h);
            }
        }
        // 清理资源
        filterBoxes.clear();
        objProbs.clear();
        classId.clear();
        indexArray.clear();
        printf("detect done\n");
        RK_MPI_SYS_SendMediaBuffer(RK_ID_VENC, 0, mb);
        printf("\n---------------------------\n");
        RK_MPI_MB_ReleaseBuffer(mb);
        RK_MPI_MB_ReleaseBuffer(MB);
    }

    free(ctx.out_data);
    free(ctx.input_attrs);
    free(ctx.output_attrs);
}